"""
Bounding Volume Hierarchies
"""
import numpy as np
import taichi as ti

from pyrender.gmath import normalized,identity_mat4
from .scenes import (
    Scenes,
    Renderable
    )
from .gtaichi import * # pylint: disable=wildcard-import,unused-wildcard-import

TEX_WIDTH = 64*64

@ti.dataclass
class BvhNode:
    """
    Bvh二叉树节点
    """
    min     :vec3
    max     :vec3  #aabb包围盒
    left    :ti.u32
    right   :ti.u32

@ti.dataclass
class TrangleInd:
    """
    三角形顶点a,b,c,法线方向n
    material 材质索引
    """
    a : ti.u32
    b : ti.u32
    c : ti.u32
    a_uv : ti.u32
    b_uv : ti.u32
    c_uv : ti.u32
    n : ti.u32 #normal
    m : ti.u32 #material

@ti.data_oriented
class Bvh:
    """
    场景加速结构
    """
    # pylint: disable=no-member
    def __init__(self,scene):
        self.scene = scene
        self.obj2tex = {}
        self.obj2material = {}
        self.texname_tex = {}
        self.alloc()
        self.load_texture()
        self.load_material()
        self.load_primitive()

    def alloc(self):
        """
        重新分配全部数据
        """
        obj_n,tran_n,tex_n,sphere_n = self.get_scene_info()
        aabb_s = ti.root.dynamic(ti.i,1024,chunk_size=32)
        self.root = ti.field(BvhNode)
        aabb_s.place(self.root)
        self.sphere = SphereBound.field(shape=(sphere_n+1,))
        self.material = Material.field(shape=(obj_n+1,))
        self.trangle = TrangleInd.field(shape=(tran_n+1,))
        self.vertex = ti.Vector.field(3,dtype=ti.f32,shape=(tran_n+1,))
        self.normal = ti.Vector.field(3,dtype=ti.f32,shape=(tran_n+1,))
        self.tex_coord = ti.Vector.field(2,dtype=ti.f32,shape=(tran_n+1,))
        self.texture = ti.Vector.field(3,dtype=ti.f32)
        ti.root.dense(ti.i,(tex_n,)).pointer(ti.jk,(64,64)).dense(ti.jk,(64,64)).place(self.texture)
        self.texture_size = ti.Vector.field(2,dtype=ti.i32,shape=(tex_n,))

    def load_material(self):
        """
        重新加载材质
        """
        i = 0
        for obj in self.scene.objs:
            m = obj.material
            tex = self.obj2tex[id(obj)]
            t = MaterialType.BRDF
            color = obj.color
            a = m[0]
            if obj.light is not None:
                color = obj.light[:3]
                t = obj.light[3]
                a = obj.light[4]
            self.material[i] = Material(t,color,tex,a,m[1],m[2],m[3],m[4])
            self.obj2material[id(obj)] = i
            i += 1

    def load_texture(self):
        """
        重新加载材质贴图
        """
        tex = 0
        for obj in self.scene.objs:
            if obj.tex_name:
                b,ins = self.load_tex(obj.tex_name,tex)
                if b:
                    self.obj2tex[id(obj)] = tex
                    tex += ins

    @ti.kernel
    def copy_tex(self,tex:int,img:ti.types.ndarray()):# pylint: disable=no-value-for-parameter
        """
        将numpy数组图像复制到对应的纹理,剪切掉大于TEX_WIDTH的部分
        """
        for i in range(min(TEX_WIDTH,img.shape[0])):
            for j in range(min(TEX_WIDTH,img.shape[1])):
                self.texture[tex,i,j] = ti.Vector([img[i,j,0]/255.,
                                                    img[i,j,1]/255.,
                                                    img[i,j,2]/255.])

    def load_tex(self,name,tex):
        """
        加载纹理图name,返回纹理索引
        """
        if name in self.texname_tex:
            return self.texname_tex[name],0
        if name:
            img = ti.tools.imread(name)
            if img is not None:
                self.texture_size[tex][0] = img.shape[0]
                self.texture_size[tex][1] = img.shape[1]
                img_2 = np.ndarray(dtype=np.uint8,shape=img.shape)
                img_2[:] = img[:]
                self.copy_tex(tex,img_2)
                self.texname_tex[name] = tex
                return tex,1
        return -1,0

    def load_primitive(self):
        """
        加载顶点,法线,三角形,贴图坐标
        """
        sphere_i,vex_i,tran_i,tex_i = 0,0,0,0
        np_vex = np.ndarray((self.vertex.shape[0],3),dtype=np.float32)
        np_norm = np.ndarray((self.vertex.shape[0],3),dtype=np.float32)
        np_tex = np.ndarray((self.vertex.shape[0],2),dtype=np.float32)
        np_tran = np.ndarray((self.trangle.shape[0],8),dtype=np.uint32)
        def proccess_geometry(mat,o):
            mesh = o.mesh
            if mesh['type'] == 'sphere':
                center = mat @ np.array([0.,0.,0.,1.])
                x_axis = mat @ np.array([1.,0.,0.,1.])
                y_axis = mat @ np.array([0.,1.,0.,1.])
                z_axis = mat @ np.array([0.,0.,1.,1.])
                self.sphere[sphere_i] = SphereBound(0,
                                                    self.obj2material[id(o)],
                                                    ti.Vector(center[:3]),
                                                    mesh['meta'][0],
                                                    ti.Vector(x_axis[:3]),
                                                    ti.Vector(y_axis[:3]),
                                                    ti.Vector(z_axis[:3]))
            else:
                points = (mat @ mesh['points'].T).T
                trangle = mesh['trangles']
                tex_coord = mesh['textures']
                #np_tran[] = 
                np_vex[vex_i:vex_i+points.shape[0],:] = points[:,:3]
                np_tex[tex_i:tex_i+tex_coord.shape[0],:] = tex_coord[:,:2]
                tex_i += tex_coord.shape[0]
                vex_i += points.shape[0]

        self.scene.for_each(proccess_geometry,identity_mat4())

    def build(self):
        """
        创建bvh结构
        """
                
    def get_scene_info(self):
        """
        取得场景中的对象数量,三角形数量,材质数量
        """
        obj_n,tran_n,tex_n,sphere_n = 0,0,0,0
        obj_n = len(self.scene.objs)
        for obj in self.scene.objs:
            if obj.tex_name:
                tex_n += 1
            if obj.mesh and 'trangles' in obj.mesh:
                tran_n += obj.mesh['trangles'].shape[0]
                if obj.mesh['type']=='sphere':
                    sphere_n += 1
        return obj_n,tran_n,tex_n,sphere_n
    
    def ray_intersection(self,ray:Ray):
        pass
